#include "syncbntrainingupdate_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_syncbntrainingupdate_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void syncbntrainingupdate_custom(GM_ADDR local_mean, GM_ADDR local_var,
    GM_ADDR global_mean, GM_ADDR global_var,
    GM_ADDR updated_mean, GM_ADDR updated_var,
    SyncBNTrainingUpdateCustomTilingData tiling);
#endif

extern void GenerateTilingData(uint8_t* tilingBuf, uint32_t blockDim);

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;  // 可根据实际需求调整block维度
    size_t tilingSize = sizeof(SyncBNTrainingUpdateCustomTilingData);  // 同步BN的tiling数据大小
    uint8_t *tiling = nullptr;

#ifdef ASCENDC_CPU_DEBUG
    // CPU调试模式下的内存分配
    tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
 #else   
    // 设备模式初始化
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    // 定义主机和设备端指针
    uint8_t *localMeanHost, *localVarHost, *globalMeanHost, *globalVarHost;
    uint8_t *updatedMeanHost, *updatedVarHost;
    uint8_t *localMeanDevice, *localVarDevice, *globalMeanDevice, *globalVarDevice;
    uint8_t *updatedMeanDevice, *updatedVarDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
#endif
    GenerateTilingData(tiling, blockDim);
    uint32_t totalLength = reinterpret_cast<SyncBNTrainingUpdateCustomTilingData *>(tiling)->totalLength;
    size_t inputByteSize = totalLength * sizeof(float);// 应为16384*4=65536字节
    size_t outputByteSize = totalLength * sizeof(float);
#ifdef ASCENDC_CPU_DEBUG
    // 分配输入输出全局内存
    uint8_t *localMean = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *localVar = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *globalMean = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *globalVar = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *updatedMean = (uint8_t *)AscendC::GmAlloc(outputByteSize);
    uint8_t *updatedVar = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    // 读取输入数据
    ReadFile("./input/input_local_mean.bin", inputByteSize, localMean, inputByteSize);
    ReadFile("./input/input_local_var.bin", inputByteSize, localVar, inputByteSize);
    ReadFile("./input/input_global_mean.bin", inputByteSize, globalMean, inputByteSize);
    ReadFile("./input/input_global_var.bin", inputByteSize, globalVar, inputByteSize);

    // 设置内核模式并运行
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(syncbntrainingupdate_custom, blockDim, 
                localMean, localVar, globalMean, globalVar, 
                updatedMean, updatedVar,
                *reinterpret_cast<SyncBNTrainingUpdateCustomTilingData *>(tiling));

    // 写入输出结果
    WriteFile("./output/output_updated_mean.bin", updatedMean, outputByteSize);
    WriteFile("./output/output_updated_var.bin", updatedVar, outputByteSize);

    // 释放内存
    AscendC::GmFree((void *)localMean);
    AscendC::GmFree((void *)localVar);
    AscendC::GmFree((void *)globalMean);
    AscendC::GmFree((void *)globalVar);
    AscendC::GmFree((void *)updatedMean);
    AscendC::GmFree((void *)updatedVar);
    AscendC::GmFree((void *)tiling);
#else

    // 分配主机端内存
    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    CHECK_ACL(aclrtMallocHost((void **)(&localMeanHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&localVarHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&globalMeanHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&globalVarHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&updatedMeanHost), outputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&updatedVarHost), outputByteSize));

    // 分配设备端内存
    CHECK_ACL(aclrtMalloc((void **)&localMeanDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&localVarDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&globalMeanDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&globalVarDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&updatedMeanDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&updatedVarDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // 读取输入数据到主机端
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
    ReadFile("./input/input_local_mean.bin", inputByteSize, localMeanHost, inputByteSize);
    ReadFile("./input/input_local_var.bin", inputByteSize, localVarHost, inputByteSize);
    ReadFile("./input/input_global_mean.bin", inputByteSize, globalMeanHost, inputByteSize);
    ReadFile("./input/input_global_var.bin", inputByteSize, globalVarHost, inputByteSize);

    // 主机到设备数据拷贝
    CHECK_ACL(aclrtMemcpy(localMeanDevice, inputByteSize, localMeanHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(localVarDevice, inputByteSize, localVarHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(globalMeanDevice, inputByteSize, globalMeanHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(globalVarDevice, inputByteSize, globalVarHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // 启动内核
    ACLRT_LAUNCH_KERNEL(syncbntrainingupdate_custom)(blockDim, stream, 
                                                     localMeanDevice, localVarDevice, 
                                                     globalMeanDevice, globalVarDevice, 
                                                     updatedMeanDevice, updatedVarDevice, 
                                                    reinterpret_cast<SyncBNTrainingUpdateCustomTilingData *>tiling);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    // 设备到主机数据拷贝
    CHECK_ACL(aclrtMemcpy(updatedMeanHost, outputByteSize, updatedMeanDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    CHECK_ACL(aclrtMemcpy(updatedVarHost, outputByteSize, updatedVarDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));

    // 写入输出结果
    WriteFile("./output/output_updated_mean.bin", updatedMeanHost, outputByteSize);
    WriteFile("./output/output_updated_var.bin", updatedVarHost, outputByteSize);

    // 释放设备内存
    CHECK_ACL(aclrtFree(localMeanDevice));
    CHECK_ACL(aclrtFree(localVarDevice));
    CHECK_ACL(aclrtFree(globalMeanDevice));
    CHECK_ACL(aclrtFree(globalVarDevice));
    CHECK_ACL(aclrtFree(updatedMeanDevice));
    CHECK_ACL(aclrtFree(updatedVarDevice));
    // 释放主机内存
    CHECK_ACL(aclrtFreeHost(localMeanHost));
    CHECK_ACL(aclrtFreeHost(localVarHost));
    CHECK_ACL(aclrtFreeHost(globalMeanHost));
    CHECK_ACL(aclrtFreeHost(globalVarHost));
    CHECK_ACL(aclrtFreeHost(updatedMeanHost));
    CHECK_ACL(aclrtFreeHost(updatedVarHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    // 清理资源
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}